{
 "cells": [
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2024-12-01T08:01:27.899479Z",
     "start_time": "2024-12-01T08:01:20.593965Z"
    }
   },
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torchvision\n",
    "from torchvision import transforms\n",
    "import numpy as np"
   ],
   "outputs": [],
   "execution_count": 1
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-01T08:01:31.972044Z",
     "start_time": "2024-12-01T08:01:31.763044Z"
    }
   },
   "cell_type": "code",
   "source": [
    "train_dataset = torchvision.datasets.MNIST(root='./data', train=True, transform=transforms.ToTensor())\n",
    "test_dataset = torchvision.datasets.MNIST(root='./data', train=False, transform=transforms.ToTensor())"
   ],
   "id": "2ca3c9168dafc9ac",
   "outputs": [],
   "execution_count": 2
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-01T08:01:32.575833Z",
     "start_time": "2024-12-01T08:01:32.568840Z"
    }
   },
   "cell_type": "code",
   "source": [
    "train_loader=torch.utils.data.DataLoader(dataset=train_dataset, batch_size=64, shuffle=True)\n",
    "test_loader=torch.utils.data.DataLoader(dataset=test_dataset, batch_size=64, shuffle=True)"
   ],
   "id": "ab04788080621ae",
   "outputs": [],
   "execution_count": 3
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-01T08:01:33.403543Z",
     "start_time": "2024-12-01T08:01:33.357041Z"
    }
   },
   "cell_type": "code",
   "source": [
    "for x, y in train_loader:\n",
    "    #print(y)\n",
    "    print(x.view(x.shape[0], -1).shape)\n",
    "    #print(np.array(x.view(x.shape[0],-1)[0]).shape)\n",
    "    \n",
    "    break"
   ],
   "id": "cc4be81bb4de1169",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([64, 784])\n"
     ]
    }
   ],
   "execution_count": 4
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "原版",
   "id": "c943ab24127ab064"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-01T08:02:25.049333Z",
     "start_time": "2024-12-01T08:02:25.043653Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class Net(nn.Module):\n",
    "    def __init__(self,hidden_size=128,n_classes=10):\n",
    "        super(Net, self).__init__()\n",
    "        self.fc1 = nn.Linear(in_features=28*28, out_features=hidden_size)\n",
    "        self.fc2 = nn.Linear(in_features=hidden_size, out_features=n_classes)\n",
    "        self.softmax = nn.LogSoftmax(dim=1)\n",
    "        self.relu = nn.ReLU()\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.relu(self.fc1(x))\n",
    "        x = self.relu(self.fc2(x))\n",
    "        x = self.softmax(x)\n",
    "        return x\n",
    "    \n",
    "    def predict(self,x,y):\n",
    "        out = self.forward(x)\n",
    "        loss = criterion(out, y)\n",
    "        print(loss.item())\n",
    "        "
   ],
   "id": "c6516ab0c1342f14",
   "outputs": [],
   "execution_count": 9
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "手动实现",
   "id": "93a7eb8aec5554e0"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-01T09:49:27.943377Z",
     "start_time": "2024-12-01T09:49:27.932767Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class NetH(nn.Module):\n",
    "    def __init__(self,hidden_size=128,n_classes=10):\n",
    "        super(NetH, self).__init__()\n",
    "        self.fc1 = nn.Linear(in_features=28*28, out_features=hidden_size)\n",
    "        self.fc2 = nn.Linear(in_features=hidden_size, out_features=n_classes)\n",
    "        self.softmax = nn.LogSoftmax(dim=1)\n",
    "        self.relu = nn.ReLU()\n",
    "        \n",
    "    def forward(self, x):\n",
    "        \n",
    "        x = self.relu(self.fc1(x))\n",
    "        x = self.relu(self.fc2(x))\n",
    "        x = self.softmax(x)\n",
    "        return x\n",
    "    \n",
    "    def predict(self,x,y):\n",
    "        out = self.forward(x)\n",
    "        loss = criterion(out, y)\n",
    "        print(loss.item())\n",
    "        "
   ],
   "id": "4ce88ef177a1ba67",
   "outputs": [],
   "execution_count": 71
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "torch.nn实现",
   "id": "8edcd0570a5628f6"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-01T09:54:03.781003Z",
     "start_time": "2024-12-01T09:54:03.773164Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class NetM(nn.Module):\n",
    "    def __init__(self,hidden_size=128,n_classes=10):\n",
    "        super(NetM, self).__init__()\n",
    "        self.fc1 = nn.Linear(in_features=28*28, out_features=hidden_size)\n",
    "        self.fc2 = nn.Linear(in_features=hidden_size, out_features=n_classes)\n",
    "        self.dropout=nn.Dropout(0.5)\n",
    "        self.softmax = nn.LogSoftmax(dim=1)\n",
    "        self.relu = nn.ReLU()\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.relu(self.fc1(x))\n",
    "        x=self.dropout(x)\n",
    "        x = self.relu(self.fc2(x))\n",
    "        x = self.softmax(x)\n",
    "        return x\n",
    "    \n",
    "    def predict(self,x,y):\n",
    "        out = self.forward(x)\n",
    "        loss = criterion(out, y)\n",
    "        print(loss.item())\n",
    "        "
   ],
   "id": "6a433199ba2734fd",
   "outputs": [],
   "execution_count": 76
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-01T08:24:33.112453Z",
     "start_time": "2024-12-01T08:24:33.103089Z"
    }
   },
   "cell_type": "code",
   "source": [
    "model = Net()\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.001)"
   ],
   "id": "673eaed091ca754",
   "outputs": [],
   "execution_count": 20
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-01T09:49:31.776049Z",
     "start_time": "2024-12-01T09:49:31.766281Z"
    }
   },
   "cell_type": "code",
   "source": [
    "modelH = NetH()\n",
    "optimizerH = torch.optim.Adam(modelH.parameters(), lr=0.001)\n",
    "np.random.seed(0)"
   ],
   "id": "69edb4b09efd120f",
   "outputs": [],
   "execution_count": 72
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-01T09:54:06.763914Z",
     "start_time": "2024-12-01T09:54:06.757620Z"
    }
   },
   "cell_type": "code",
   "source": [
    "modelM = NetH()\n",
    "optimizerM = torch.optim.Adam(modelM.parameters(), lr=0.001)"
   ],
   "id": "52a7898a32205f81",
   "outputs": [],
   "execution_count": 77
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-01T09:54:10.626748Z",
     "start_time": "2024-12-01T09:54:10.621557Z"
    }
   },
   "cell_type": "code",
   "source": [
    "def train(model, train_loader, optimizer,epochs):#torch与原版共享同一个训练函数\n",
    "    model.train()\n",
    "    for epoch in range(epochs):\n",
    "        for x, y in train_loader:\n",
    "            x=x.view(x.shape[0], -1) \n",
    "            output = model(x)\n",
    "            loss = criterion(output, y)\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "        print('Epoch: {}, Loss: {}'.format(epoch, loss.item()))"
   ],
   "id": "b2b4802854f2aa00",
   "outputs": [],
   "execution_count": 78
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-01T09:50:21.660493Z",
     "start_time": "2024-12-01T09:50:21.652904Z"
    }
   },
   "cell_type": "code",
   "source": [
    "def trainH(model, train_loader, optimizer,epochs):\n",
    "    model.train()\n",
    "    for epoch in range(epochs):\n",
    "        mask=np.random.choice([0, 1], size=(1, 784))#才疏学浅，本想在隐藏层加mask，但是会丢gradient，所以擒贼擒王了:)\n",
    "        mask = mask.astype(float)\n",
    "        mask= torch.from_numpy(mask)\n",
    "        for x, y in train_loader:\n",
    "            x=x.view(x.shape[0], -1)\n",
    "            for it in range(x.shape[0]):\n",
    "                x[it]=torch.mul(x[it], mask)\n",
    "            output = model(x)\n",
    "            loss = criterion(output, y)\n",
    "            optimizerH.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizerH.step()\n",
    "        print('Epoch: {}, Loss: {}'.format(epoch, loss.item()))"
   ],
   "id": "11740398004891c3",
   "outputs": [],
   "execution_count": 73
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "原版结果",
   "id": "707363712ffd7102"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-01T08:27:49.303531Z",
     "start_time": "2024-12-01T08:25:43.544888Z"
    }
   },
   "cell_type": "code",
   "source": "train(model,train_loader,optimizer,10)",
   "id": "7a61d919e6623265",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 0, Loss: 0.16410426795482635\n",
      "Epoch: 1, Loss: 0.0878652036190033\n",
      "Epoch: 2, Loss: 0.040897827595472336\n",
      "Epoch: 3, Loss: 0.1996774673461914\n",
      "Epoch: 4, Loss: 0.005071149207651615\n",
      "Epoch: 5, Loss: 0.03542032092809677\n",
      "Epoch: 6, Loss: 0.052408792078495026\n",
      "Epoch: 7, Loss: 0.01362878829240799\n",
      "Epoch: 8, Loss: 0.007548083551228046\n",
      "Epoch: 9, Loss: 0.008145295083522797\n"
     ]
    }
   ],
   "execution_count": 25
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "手工实现结果",
   "id": "a43329e033a54b6d"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-01T09:53:52.292680Z",
     "start_time": "2024-12-01T09:51:23.425792Z"
    }
   },
   "cell_type": "code",
   "source": "trainH(modelH,train_loader,optimizerH,10)",
   "id": "8d0cf47922abed7e",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 0, Loss: 0.3721144199371338\n",
      "Epoch: 1, Loss: 0.2991145849227905\n",
      "Epoch: 2, Loss: 0.259930819272995\n",
      "Epoch: 3, Loss: 0.16070783138275146\n",
      "Epoch: 4, Loss: 0.21261680126190186\n",
      "Epoch: 5, Loss: 0.30492281913757324\n",
      "Epoch: 6, Loss: 0.5076764822006226\n",
      "Epoch: 7, Loss: 0.41066765785217285\n",
      "Epoch: 8, Loss: 0.3833639621734619\n",
      "Epoch: 9, Loss: 0.2573384940624237\n"
     ]
    }
   ],
   "execution_count": 75
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "torch.nn实现结果",
   "id": "b7290673cb4790f7"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-01T09:56:53.387473Z",
     "start_time": "2024-12-01T09:54:49.653479Z"
    }
   },
   "cell_type": "code",
   "source": "train(modelM,train_loader,optimizerM,10)",
   "id": "7a35acbee5b23877",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 0, Loss: 0.6608747839927673\n",
      "Epoch: 1, Loss: 0.43181395530700684\n",
      "Epoch: 2, Loss: 0.2906467318534851\n",
      "Epoch: 3, Loss: 0.33798205852508545\n",
      "Epoch: 4, Loss: 0.298394113779068\n",
      "Epoch: 5, Loss: 0.26280057430267334\n",
      "Epoch: 6, Loss: 0.22039049863815308\n",
      "Epoch: 7, Loss: 0.41757097840309143\n",
      "Epoch: 8, Loss: 0.22273992002010345\n",
      "Epoch: 9, Loss: 0.30578556656837463\n"
     ]
    }
   ],
   "execution_count": 80
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": "",
   "id": "b0d842d1bcb669a3"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
